import cv2
import numpy as np
import os
import sys
from PIL import Image

def detect_circular_plates(image_path, output_dir, min_radius=50, max_radius=200):
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 读取图片并转换为灰度图
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"无法读取图片: {image_path}")
    
    # 预处理: 高斯模糊降噪
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (15, 15), 0)
    
    # 使用霍夫圆变换检测圆形盘子
    circles = cv2.HoughCircles(
        blurred,
        method=cv2.HOUGH_GRADIENT,
        dp=1,
        minDist=100,
        param1=50,
        param2=30,
        minRadius=min_radius,
        maxRadius=max_radius
    )
    
    if circles is None:
        print("未检测到圆形盘子")
        return
    
    # 转换为整数坐标
    circles = np.uint16(np.around(circles))
    
    # 处理每个检测到的圆
    dish_count = 0
    for i in circles[0, :]:
        center = (i[0], i[1])
        radius = i[2]
        
        # 计算裁剪区域（包含圆形盘子和下方菜名）
        name_height = int(radius * 0.03)  # 菜名区域高度（半径的3%，进一步减少下方留白）
        x1 = max(0, center[0] - radius)
        y1 = max(0, center[1] - radius)
        x2 = min(img.shape[1], center[0] + radius)
        y2 = min(img.shape[0], center[1] + radius + name_height)  # 向下扩展以包含菜名
        
        # 裁剪并保存图片
        cropped_img = img[y1:y2, x1:x2]
        output_path = os.path.join(output_dir, f"dish_{dish_count}.png")
        cv2.imwrite(output_path, cropped_img)
        print(f"已保存: {output_path}")
        dish_count += 1

if __name__ == "__main__":
    if len(sys.argv) != 5:
        print("用法: python split_dishes.py <图片路径> <最小半径> <最大半径> <输出目录>")
        print("示例: python split_dishes.py image.png 80 150 output1_improved")
        sys.exit(1)
    
    # 获取命令行参数
    image_path = sys.argv[1]
    min_radius = int(sys.argv[2])
    max_radius = int(sys.argv[3])
    output_dir = sys.argv[4]
    
    # 执行圆形盘子检测和分割
    detect_circular_plates(image_path, output_dir, min_radius, max_radius)